import os
import re
from collections import defaultdict, deque
from datetime import datetime
import threading
import time
from wxauto import WeChat
from database import DatabaseManager
from logging_config import logger
import json
from functools import lru_cache
from openai import OpenAI
from typing import List, Optional, Dict, Any
from pathlib import Path
import httpx


# 获取当前脚本所在目录
script_dir = os.path.dirname(os.path.abspath(__file__))
# 拼接配置文件的路径
config_file_path = os.path.join(script_dir, 'config.json')
logger.error(config_file_path)

# 读取 config.json 文件
try:
    with open(config_file_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
except FileNotFoundError:
    logger.error(f"配置文件 {config_file_path} 未找到")
    config = {}

# 获取配置信息
LISTEN_LIST = config['LISTEN_LIST']
CURRENT_WX_NAME = config['CURRENT_WX_NAME']
GROUP_WX_NAME = config['GROUP_WX_NAME']
DEEPSEEK_API_KEY = config['DEEPSEEK_API_KEY']
DEEPSEEK_BASE_URL = config['DEEPSEEK_BASE_URL']
MODEL = config['MODEL']
MAX_TOKEN = config['MAX_TOKEN']
TEMPERATURE = config['TEMPERATURE']
PROMPT_FILE = os.path.join(script_dir, 'prompt.md')
MAX_HISTORY_PAIRS = config['MAX_HISTORY_PAIRS']
MESSAGE_TIMEOUT = config['MESSAGE_TIMEOUT']
MAX_CONTEXT_LENGTH = config['MAX_CONTEXT_LENGTH']



class AIService:
    def __init__(self):
        self.client = OpenAI(
            api_key=DEEPSEEK_API_KEY,
            base_url=DEEPSEEK_BASE_URL
        )
        self.prompt_content = self._load_prompt()
        self.model = MODEL
        self.temperature = TEMPERATURE
        self.max_token = MAX_TOKEN
        self.prompt_cache_tag = "system_prompt_v1"

    def _load_prompt(self):
        """加载提示词文件"""
        try:
            with open(PROMPT_FILE, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            logger.error(f"提示词加载失败：{str(e)}")
            return ""

    def upload_files(self, files: List[str], cache_tag: Optional[str] = None) -> List[Dict[str, Any]]:
        """（保持与示例完全一致的实现）"""
        messages = []

        for file in files:
            file_object = self.client.files.create(file=Path(file), purpose="file-extract")
            file_content = self.client.files.content(file_id=file_object.id).text
            messages.append({
                "role": "system",
                "content": file_content,
            })

        if cache_tag:
            r = httpx.post(f"{self.client.base_url}caching",
                           headers={"Authorization": f"Bearer {self.client.api_key}"},
                           json={
                               "model": "moonshot-v1",
                               "messages": messages,
                               "ttl": 300,
                               "tags": [cache_tag],
                           })
            logger.error(f"缓存上传结果：{r.text}")
            if r.status_code != 200:
                raise Exception(r.text)
            return [{
                "role": "cache",
                "content": f"tag={cache_tag};reset_ttl=300",
            }]
        return messages

    @lru_cache(maxsize=128)
    def get_response(self, user_id, context_json):
        """获取AI回复，并记录上下文内容"""
        context = json.loads(context_json)
        # 记录请求开始时间
        start_time = time.time()

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    *self._get_cache_messages(),
                    *context
                ],
                temperature=self.temperature,
                max_tokens=self.max_token
            )
            # 记录 token 使用量
            prompt_tokens = response.usage.prompt_tokens
            completion_tokens = response.usage.completion_tokens
            # 记录请求结束时间
            end_time = time.time()
            elapsed_time = end_time - start_time
            logger.error(f"Token 使用量 - 用户：{user_id}，输入：{prompt_tokens}，输出：{completion_tokens}，耗时：{elapsed_time:.2f} 秒")
            return response.choices[0].message.content.strip()
        except Exception as e:
            logger.error(f"API调用失败：{str(e)}，user_id: {user_id}，messages: {context}")
            return "睡着了..."

    def _get_cache_messages(self):
        """获取缓存消息"""
        if not hasattr(self, '_cache_messages'):
            self._cache_messages = []
            if os.path.exists(PROMPT_FILE):
                try:
                    self._cache_messages = self.upload_files(
                        files=[PROMPT_FILE],
                        cache_tag=self.prompt_cache_tag
                    )
                except Exception as e:
                    logger.error(f"缓存初始化失败: {str(e)}")
        return self._cache_messages


class WeChatBot:
    def __init__(self):
        self.wx = WeChat()
        self.ai = AIService()
        self.user_queues = {}
        self.chat_contexts = defaultdict(lambda: deque(maxlen=MAX_CONTEXT_LENGTH))
        self.lock = threading.Lock()
        self.current_listeners = set()
        self.message_timeout = MESSAGE_TIMEOUT
        self._setup_listening()
        self._load_initial_contexts()
        self.message_lock = threading.Lock()


    def _setup_listening(self):
        """初始化监听列表"""
        if not all(isinstance(item, str) for item in LISTEN_LIST):
            raise ValueError("LISTEN_LIST 中包含不可哈希的元素")
        new_listeners = set(LISTEN_LIST) - self.current_listeners
        removed_listeners = self.current_listeners - set(LISTEN_LIST)

        for name in new_listeners:
            try:
                self.wx.AddListenChat(who=name, savepic=False)
                logger.info(f"已添加监听对象：{name}")
            except Exception as e:
                logger.error(f"添加监听对象 {name} 失败：{str(e)}")

        for name in removed_listeners:
            try:
                self.wx.RemoveListenChat(who=name)
                logger.info(f"已移除监听对象：{name}")
            except Exception as e:
                logger.error(f"移除监听对象 {name} 失败：{str(e)}")

        self.current_listeners = set(LISTEN_LIST)

    def _load_initial_contexts(self):
        """初始化时加载上下文"""
        with self.lock:
            for name in LISTEN_LIST:
                recent_msgs = DatabaseManager.get_recent_messages(name)
                for msg in recent_msgs:
                    self.chat_contexts[name].extend([
                        {"role": "user", "content": msg["message"]},
                        {"role": "assistant", "content": msg["reply"]}
                    ])

    def _process_message(self, is_group=False, **kwargs):
        msg = kwargs['msg']
        content = getattr(msg, 'content', '')
        sender_id = msg.sender if not is_group else kwargs['chat_id']
        sender_name = msg.sender if not is_group else kwargs['chat_name']

        timestamp_msg = f"[{datetime.now():%Y-%m-%d %H:%M:%S}] {content}"

        self._add_message_to_queue(sender_id, timestamp_msg)
        logger.info(f"消息已排队：{sender_name}")

    def _add_message_to_queue(self, sender_id, message):
        with self.lock:
            if sender_id not in self.user_queues:
                self.user_queues[sender_id] = {
                    'timer': threading.Timer(
                        self.message_timeout,
                        self._handle_message_queue,
                        args=(sender_id,)
                    ),
                    'messages': [message],
                    'context_key': sender_id
                }
                self.user_queues[sender_id]['timer'].start()
            else:
                self.user_queues[sender_id]['messages'].append(message)
                self.user_queues[sender_id]['timer'].cancel()  # 取消旧计时器
                self.user_queues[sender_id]['timer'] = threading.Timer(
                    self.message_timeout,
                    self._handle_message_queue,
                    args=(sender_id,)
                )
                self.user_queues[sender_id]['timer'].start()  # 启动新计时器

    def _handle_message_queue(self, user_id):
        with self.lock:
            if user_id not in self.user_queues:
                return
            data = self.user_queues.pop(user_id)

        merged = ' \\ '.join(data['messages'])
        merged = merged.encode('gbk', 'ignore').decode('gbk')  # 过滤无法编码的字符
        logger.info(f"处理合并消息：{merged}")

        context = self._get_context(user_id)
        temp_context = context + [{"role": "user", "content": merged}]

        context_json = json.dumps(temp_context[-MAX_CONTEXT_LENGTH:])
        reply = self.ai.get_response(user_id, context_json)

        self._update_context(user_id, merged, reply, keep_recent=MAX_HISTORY_PAIRS*2)

        self._send_reply(reply, user_id)

        DatabaseManager.save_message(
            sender_id=user_id,
            sender_name=user_id,
            message=merged,
            reply=reply
        )

    def _get_recent_context(self, user_id):
        """获取最近的N组对话，而非全部历史"""
        return list(self.chat_contexts[user_id])[-MAX_HISTORY_PAIRS * 2:]

    def _get_context(self, user_id):
        return list(self.chat_contexts[user_id])

    def _update_context(self, user_id, user_message, reply, keep_recent=4):
        with self.lock:
            # 保留最近的N组对话
            recent = list(self.chat_contexts[user_id])[-keep_recent:]
            self.chat_contexts[user_id] = deque(recent, maxlen=MAX_CONTEXT_LENGTH)

            # 添加新对话
            self.chat_contexts[user_id].extend([
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": reply}
            ])

    def _send_reply(self, reply, user_id):
        try:
            for part in reply.split('\\'):
                clean_part = part.strip()
                if clean_part:
                    self.wx.SendMsg(clean_part, user_id)
                    time.sleep(1)
        except Exception as e:
            logger.error(f"消息发送失败：{str(e)}")

    def message_listener(self):
        """消息监听主循环"""
        prev_is_system = False  # 新增状态标记
        while True:
            try:
                with self.message_lock:
                    msgs = self.wx.GetListenMessage()
                for chat in msgs:
                    who = chat.who
                    one_msgs = msgs.get(chat)
                    for i, msg in enumerate(one_msgs):
                        # 新增消息顺序判断逻辑
                        if i > 0 and prev_is_system:
                            prev_msg = one_msgs[i - 1]
                            # 验证前序消息是系统消息且包含指定内容
                            if (prev_msg.sender == 'SYS'
                                    and '以下为新消息' in prev_msg.content):
                                # logger.info(f'发现有效消息序列：{prev_msg.content} -> {msg.content}')
                                prev_is_system = False  # 重置标记
                            else:
                                continue  # 不满足条件则跳过
                        elif msg.sender == 'SYS' and '以下为新消息' in msg.content:
                            # 检测系统消息并设置标记
                                prev_is_system = True
                                continue  # 跳过系统消息本身处理
                        else:
                            print(f"非最新聊天消息")
                            continue

                        # 原有消息处理逻辑
                        msgtype = msg.type
                        content = msg.content
                        print(f'【{who}】：{content}')
                        if msgtype != 'friend':
                            print(f"非聊天消息，忽略: {msgtype}")
                            continue

                        if chat.who == msg.sender:
                            self._process_message(msg=msg)
                        else:
                            if re.search(f'@{GROUP_WX_NAME}', msg.content):
                                self._process_message(
                                    is_group=True,
                                    msg=msg,
                                    chat_id=chat.who,
                                    chat_name=chat.who,
                                    mentioned_user=msg.sender
                                )
                    time.sleep(1)
            except Exception as e:
                logger.error(f"监听异常：{str(e)}")

    def add_to_listen_list(self, nickname):
        """添加好友到监听列表"""
        if nickname not in LISTEN_LIST:
            LISTEN_LIST.append(nickname)
            # 更新配置文件
            self._update_config()
            # 重新初始化监听列表
            self._setup_listening()

    def remove_from_listen_list(self, nickname):
        """从监听列表中移除好友"""
        if nickname in LISTEN_LIST:
            LISTEN_LIST.remove(nickname)
            self._update_config()
            self._setup_listening()

    def _update_config(self):
        """更新 config.json 文件"""
        config['LISTEN_LIST'] = LISTEN_LIST
        with open('config.json', 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=4)

    def get_listen_list(self):
        # 读取配置文件
        with open('config.json', 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config['LISTEN_LIST']

    def get_chat_history(self, selected_friend):
        return DatabaseManager.get_recent_messages(selected_friend)
